# -*- coding: utf-8 -*-
"""
Native Python RPC Layer
"""
from debugtools import *
from codepickle import *
from connection import *
from multiprocessing import Process
import multiprocessing, multiprocessing.connection, threading, logging
import os, sys, zlib, cPickle, time, traceback, gc, socket, base64, math, binascii
from numpy.random import randint

__all__ = ['DEFAULT_PORT', 'RpcServer', 'RpcClient', 'RpcClients', 'server', 'close_servers']

DEFAULT_PORT = 2718




def close_servers(addresses):
    if type(addresses) is str:
        addresses = [(addresses, DEFAULTPORT)]
    if type(addresses) is tuple:
        addresses = [addresses]
    RpcClients(addresses).close_server()

def server(port = None):
    RpcServer(port=port).listen()




class Procedure(object):
    """
    A procedure sent by a client to a server.
    """
    def __init__(self, name, args, kwds, async = False):
        self.name = name
        self.args = args
        self.kwds = kwds
        self.async = async

    def __repr__(self):
        return "<Procedure %s with args %s, kwds %s>" % (self.name, str(self.args), str(self.kwds))




class DistantException(Exception):
    """
    Distant Exception class. Allows to re-raise exception on the client, giving 
    the filename and the line number of the original server-side exception. 
    """
    def __init__(self, exception = None):
        if type(exception) == str:
            exception = Exception(exception)
        self.exception = exception
        try:
            self.filename = __file__
            if os.path.splitext(__file__)[1] == '.pyc':
                self.filename = __file__[:-1]
        except:
            self.filename = 'unknown'
        try:
            self.line = sys._getframe(1).f_lineno
        except:
            self.line = 0
        try:
            self.function = sys._getframe(1).f_code.co_name
        except:
            self.function = 'unknown'
    
    def setinfo(self):
        (self.filename, self.line, self.function, self.text) = traceback.extract_tb(sys.exc_info()[2])[-1]

    def __str__(self):
        s = "A distant exception happened in:"
        s += "\n  File \"%s\", line %d, in %s" % (self.filename, self.line, str(self.function))
        s += "\n    %s" % str(self.exception)
        return s



class RpcServer(object):
    def __init__(self, handler_class = None, port = None, bindip = ''):#, *handler_args, **handler_kwds):
        """
        Server constructor.
        
        *Arguments*
        
        `handler_class`
          A class deriving from ProcedureHandler used to process the procedures
          sent to the server. There is one handler object per active client.
        """
        self.handler_class = handler_class
        self.handler_args = ()
        self.handler_kwds = dict()
        if port is None:
            port = DEFAULT_PORT
        self.port = port
        self.address = (bindip, self.port)
        self.handlers = dict() # a dictionary client_address : dict([handlers])
        self.handler_indices = dict() # for every client : the index of the current handler
    
    def exec_proc(self, handler, procedure):
#        if not isinstance(handler, self.handler_class):
#            raise DistantException("The current handler is not an instance of '%s', please create a new handler using a different client name" % (self.handler_class.__name__))
        try:
            log.debug("server: running procedure '%s'" % (procedure.name))
            result = getattr(handler, procedure.name)(*procedure.args, **procedure.kwds)
        except Exception as exc:
            log.warn("server: an exception occurred while running procedure '%s(*%s, **%s)'" % (procedure.name, str(procedure.args), str(procedure.kwds)))
            result = DistantException(exc)
            result.setinfo()
        return result
    
    def _connect(self, client, procedure):
        return True
    
    def _disconnect(self, client, procedure):
        return True
    
    def _close_handler(self, client, procedure):
        self.remove_handler(client)
        return True
    
    def _set_handler_class(self, client, procedure):
        self.handler_class = procedure.args[0]
        self.handler_args = procedure.args[1]
        self.handler_kwds = procedure.args[2]
        log.debug("server: handler class set to '%s'" % procedure.args[0].__name__)
        return True
    
    def _set_handler_index(self, client, procedure):
        # sets the current handler index for the current client
        self.handler_indices[client] = procedure.args[0]
        log.debug("server: handler index set to '%s'" % procedure.args[0])
        return True
    
    def _get_current_handler_index(self, client, procedure):
        return self.handler_indices[client]
    
    def _get_next_handler_index(self, client, procedure):
        if client not in self.handler_indices.keys():
            return 0
        elif self.handler_indices[client] == 0:
            if 0 in self.handlers[client].keys() and self.handlers[client][0] is not None:
                return 1
            else:
                return 0
        else:
            return self.handler_indices[client]+1
    
    def get_handler(self, client):
        return self.handlers[client][self.handler_indices[client]]
    
    def dispatch(self, client, procedure):
        """
        Dispatches the procedure to the handler, by calling handler.procedure_name(args).
        
        *Arguments*
        
        `procedure`
          A Procedure object.
        """
        result = None
        if procedure.name[0] == '_':
            # procedures starting with '_': call them directly
            if not hasattr(self, procedure.name):
                result = DistantException("The RpcServer class doesn't have a '%s' method" % procedure.name)
            else:
                result = getattr(self, procedure.name)(client, procedure)
        else:
            # any non-reserved procedure
            if client not in self.handlers or self.handler_indices[client] not in self.handlers[client].keys():
                self.add_handler(client)
            handler = self.get_handler(client)
            if hasattr(handler, procedure.name):
                result = self.exec_proc(handler, procedure)
            else:
                result = DistantException("Procedure '%s' not implemented by the handler." % procedure.name)
        return result
    
    def update_handlers(self):
        for key in self.handlers.iterkeys():
            for i in xrange(len(self.handlers[key])):
                self.handlers[key][i].handlers = self.handlers
    
    def add_handler(self, client):
        """
        Creates a new handler associated to the given client, if it doesn't
        exist yet.
        """
        if client not in self.handlers.keys():
            # creates a new handler for the client if it doesn't exist yet
            self.handlers[client] = dict()
            self.handler_indices[client] = 0
        
        # Creates a new handler at the current handler index
        handler = self.handler_class(*self.handler_args, **self.handler_kwds)
        handler.client = client
        self.handlers[client][self.handler_indices[client]] = handler
        log.debug("server: handler '%s' added for client '%s'" % (self.handler_class.__name__, client))
        # updates the handlers property of the handlers, so that each handler can access the other handlers
        self.update_handlers()
    
    def remove_handler(self, client):
        """
        Closes the handler associated to the given client.
        """
        if client in self.handlers.keys():
            handler = self.get_handler(client)
            if hasattr(handler, '__close__'):
                getattr(handler, '__close__')()
            log.debug("server: handler '%s' deleted for client '%s'" % (self.handler_class.__name__, client))
#            del handler
            del self.handlers[client][self.handler_indices[client]]
            self.handler_indices[client] = max(0, self.handler_indices[client]-1)
        # updates the handlers property of the handlers, so that
        # each handler can access the other handlers
        self.update_handlers()
    
    def _serve(self, conn, client, procedure, index = 0):
        stay_connected = None # allows to receive several procedures during the same session
        
        while stay_connected is not False:
            result = None
            
#            log.debug("NEW LOOP, procedure="+str(procedure))
            
            if procedure is None: 
#                log.debug("server: receiving procedure")
                procedure = conn.recv()
#                log.debug("server: procedure received")
            
            if procedure.name == '_connect':
                stay_connected = True                
            if procedure.name == '_disconnect':
                stay_connected = False
            if stay_connected is None:
                stay_connected = False
            
            # Dispatches the procedure to the handler unless the procedure
            # asks the server to close the handler or itself
            log.debug("server: dispatching procedure '%s'" % procedure.name)
            result = self.dispatch(client, procedure)
                
            # returns the result to the client
            log.debug("server: returning the result to the client")
            conn.send(result)
            
            procedure = None
        
        if conn is not None:
            conn.close()
            del conn
        log.debug("server: connection closed")
    
    def listen(self):
        """
        Listens to incoming connections and create one handler for each new connection.
        """
        conn = None
        
        threads = []
        index = 0
        while True:
            try:
                log.debug("server: waiting for incoming connection on port %d..." % self.address[1])
                conn, client = accept(self.address) # accepts an incoming connection
                log.info("server: connection established with client '%s' on port %d" % (client, self.address[1]))
            except:
                log.warn("server: connection NOT established, closing now")
                break         
            
            procedure = conn.recv()
            if procedure is None:
                break
            
            thread = threading.Thread(target=self._serve, args=(conn, client, procedure, index))
            thread.start()
            threads.append(thread)
            time.sleep(.05)
            index += 1

        # Closes the connection at the end of the server lifetime
        if 'conn' in locals() and conn is not None:
            conn.close()
            del conn
        
        for i in xrange(len(threads)):
            log.debug("joining thread %d" % i)
            threads[i].join()
            
        log.info("server: closing now.")


class RpcClient(object):
    """
    RPC Client constructor.
    """
    def __init__(self, server, port = None, clientid = None, handler_class = None, can_stay_connected = True):
        self.connected = False
        self.conn = None
        self.threads = []
        if type(server) is tuple:
            server, port = server
        self.server = server
        if port is None:
            port = DEFAULT_PORT
        self.port = port
        self.clientid = clientid
        if handler_class is not None:
            self.set_handler(handler_class)
    
    def _open_connection(self):
        log.debug("client: connecting to '%s' on port %d" % (self.server, self.port))
        self.conn = connect((self.server, self.port), clientid = self.clientid)
    
    def _close_connection(self):
        self.conn.close()
#        del self.conn
        self.conn = None 
        return True
    
    def _process(self, name, args, kwds):
        # Sends the procedure
        procedure = Procedure(name, args, kwds)
        
        if procedure.name == '_connect':
            self.connected = True
        elif procedure.name == '_disconnect':
            self.connected = False
            
        log.debug("client: sending procedure '%s'" % name)
        self.conn.send(procedure)
        # Receives the result
        log.debug("client: receiving result")
        try:
            result = self.conn.recv()
        except:
            result = None
            log.warn("client: no result received")
        
        if not self.connected:
            self._close_connection()
        
        return result
    
    def execute(self, name, *args, **kwds):
        """
        Calls a procedure on the server from the client.
        """
        # Connects to the server
        if not self.connected:
            self._open_connection()
        if self.conn is None:
            raise Exception("connection error")
        
        # Closes the server
        if name == '_close_server':
#            if self.conn is not None:
#                conn = self.conn
            self.conn.send(None)
            self._close_connection()
            return True
        
        # Sends the procedure to the server synchronously (blocking call on the client, but not on the server)
        result = self._process(name, args, kwds)
        
        # re-raise the Exception on the client if one was raised on the server
        if isinstance(result, Exception):
            raise result
        return result
    
    def __getattr__(self, name):
        """
        Wrapper for execute(name, *args, **kwds)
        """
        if name == 'conn':
            return None
        return lambda *args, **kwds: self.execute(name, *args, **kwds)

    def set_handler(self, handler_class, pickable = False, args = None, kwds = None):
        """
        Allows to change the handler class from the client.
        """
        if pickable:
            handler_class = PicklableClass(handler_class)
        if args is None:
            args = ()
        if kwds is None:
            kwds = dict()
        self._set_handler_class(handler_class, args, kwds)

    def add_handler(self, handler_class = None, pickable = False):
        if handler_class is not None:
            self.set_handler(handler_class, pickable)
        index = self._get_next_handler_index()
        self._set_handler_index(index)

    def connect(self):
        self._connect()
        
    def disconnect(self):
        self._disconnect()

    def close(self):
        """
        Closes the current handler associated to the client.
        """
        self._close_handler()

    def close_server(self):
        """
        Closes the server from the client.
        """
        self._close_server()
        
        

class RpcClients(object):
    def __init__(self, servers, clientid = None, handler_class = None, can_stay_connected = True):
        self.servers = servers
        self.clientid = clientid
        self.handler_class = handler_class
        
        self.clients = [RpcClient(server,
                                  clientid = clientid, 
                                  handler_class = handler_class,
                                  can_stay_connected = can_stay_connected)
                        for server in servers]
        self.results = dict()
    
    def _execute(self, client, index, name, args, kwds):
        try:
            self.results[index] = getattr(client, name)(*args, **kwds)
        except Exception as inst:
#            msg = "An exception occurred on server '%s' while executing '%s' : %s" % (client.server, name, str(inst))
            self.exceptions.append(str(inst))
    
    def distribute(self, name, *argss, **kwdss):
        """
        Makes simultaneously (multithreading) several calls to different RPC servers.
        """
        k = len(argss) # number of non-named arguments
        keys = kwdss.keys() # keyword arguments
        i = 0
        results = []
        threads = []
        self.exceptions = []
        
        # Duplicates non-list args
        argss = list(argss)
        for l in xrange(k):
            if type(argss[l]) is not list:
                argss[l] = [argss[l]]*len(self.servers)
        
        # Duplicates non-list kwds
        for key in keys:
            if type(kwdss[key]) is not list:
                kwdss[key] = [kwdss[key]]*len(self.servers)
        
        for i in xrange(len(self.servers)):
            args = [argss[l][i] for l in xrange(k)]
            kwds = dict([(key, kwdss[key][i]) for key in keys])
            thread = threading.Thread(target=self._execute, 
                                      args=(self.clients[i], i, name, args, kwds))
            thread.start()
            threads.append(thread)
        [thread.join() for thread in threads]
        if len(self.exceptions)>0:
            raise Exception("\n".join(self.exceptions))
        for i in xrange(len(self.servers)):
            if i in self.results.keys():
                result = self.results[i]
            else:
                result = None
            results.append(result)
        return results
    
    def __getattr__(self, name):
        return lambda *args, **kwds: self.distribute(name, *args, **kwds)
    

